'''
仿射加密
'''
import math
class affine_cipher():
    # 判断是否有逆元
    isReady = False
    # 保存逆元
    inverse = 0
    # 公式 f = a * x + b % m 中的a,b,m
    a = 0
    b = 0
    m = 0
    def setKey(self, a : int, b : int, m: int):
        '''
        设置 仿射加密的a,b,m
        :param a: 变量a
        :param b: 变量b
        :param m: 变量m
        :return:
        '''
        # 判断a与M是否互质
        if math.gcd(a, m) == 1:
            # 如果是互质，表示可以进行仿射加密
            self.isReady = True
            self.a = a
            self.b = b
            self.m = m
        else:
            # 不互质，不能进行加密
            self.isReady = False

    def inverse_element(self, a, m):
        '''
        计算逆元
        :param a: 变量a
        :param m: 变量m
        :return: 逆元
        '''
        m0, x0, x1 = m, 0, 1
        while a > 1:
            q = a // m
            m , a = a % m, m
            x0, x1 = x1 - q * x0, x0
        return x1 % m0


    def affine_encrypt(self, plainttext : str):
        '''
        加密算法，参照公式 f = a * x + b % m
        :param plainttext: 明文
        :return: 密文
        '''
        result = None
        if self.isReady:
            result = ""
            # 遍历每一个字符
            for x in plainttext:
                # 判断是否为字母。不是则不进行加密
                if x.isalpha():
                    # 把ASCII码转为0-25的映射数字
                    base = ord('a') if x.islower() else ord('A')
                    # 进行加密
                    result += chr(((ord(x) - base) * self.a + self.b)%self.m + base)
                else:
                    result += x
        return result

    def affine_decrypt(self, cipertext : str):
        '''
        解密算法，解密公式 f = a^-1 * (x - b) mod m
        :param cipertext: 密文
        :return: 明文
        '''
        result = None
        if self.isReady:
            result = ""
            self.inverse = self.inverse_element(self.a, self.m)
            for x in cipertext:
                # 判断是否为字母。不是则不进行加密
                if x.isalpha():
                    # 把ASCII码转为0-25的映射数字
                    base = ord('a') if x.islower() else ord('A')
                    # 进行解密
                    result += chr((self.inverse * (ord(x) - base - self.b))%self.m + base)
                else:
                    result += x
        return result